package chen.yi.component.xss.wrapper;

import cn.hutool.core.util.StrUtil;
import cn.hutool.http.HtmlUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import jakarta.servlet.ReadListener;
import jakarta.servlet.ServletInputStream;
import jakarta.servlet.http.HttpServletRequest;
import jakarta.servlet.http.HttpServletRequestWrapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.util.ContentCachingRequestWrapper;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.util.*;

/**
 * 带缓存的请求包装器
 *
 * @author huat
 **/
@Slf4j
public class CachedXssRequestWrapper extends HttpServletRequestWrapper {
    private final ContentCachingRequestWrapper cachingRequest;
    private byte[] cachedBody;
    private boolean isMultipart;

    public CachedXssRequestWrapper(HttpServletRequest request) throws IOException {
        super(request);
        this.isMultipart = isMultipartRequest(request);

        if (!isMultipart) {
            // 仅缓存非文件上传请求
            this.cachingRequest = new ContentCachingRequestWrapper(request);
            this.cachedBody = this.cachingRequest.getInputStream().readAllBytes();
            this.cachingRequest.getParameterMap(); // 触发参数缓存
        } else {
            this.cachingRequest = null;
            this.cachedBody = null;
        }
    }
    private boolean isMultipartRequest(HttpServletRequest request) {
        String contentType = request.getContentType();
        return contentType != null && contentType.startsWith("multipart/form-data");
    }
    // 清理XSS的核心方法
    private String cleanXss(String value) {
        if (value == null) return null;

        return HtmlUtil.filter(value);
    }

    // 重写输入流处理方法
    @Override
    public ServletInputStream getInputStream() throws IOException {
        //处理文件上传
        if(isMultipart){
            log.info("文件上传请求");
            return super.getRequest().getInputStream();
        }
        String param = new String(cachedBody, getCharacterEncoding());
        Object isObject = JSONUtil.parse(param);
        if (isObject instanceof JSONArray) {
            log.info("请求参数 jsonArray 类型");
            JSONArray objects = JSONUtil.parseArray(param);
            log.info("请其参数 objects:【{}】",objects);
            for (int i = 0; i < objects.size(); i++) {
                JSONObject jsonObject = objects.getJSONObject(i);
                getValue(jsonObject);
            }
            log.info("去除标签后的请求值 newObjects:【{}】",objects);
            return new CachedServletInputStream(objects.toString().getBytes(getCharacterEncoding()));
        }
        log.info("请求参数为 jsonObject 类型");
        JSONObject jsonObject = JSONUtil.parseObj(param);
        log.info(" jsonObject:【{}】",jsonObject);
        getValue(jsonObject);
        log.info("去除标签后的值 newObjects:【{}】",jsonObject);
        return new CachedServletInputStream(jsonObject.toString().getBytes(getCharacterEncoding()));
    }

    /**
     * 重写参数获取方法
     *
     * @param name
     * @return
     */
    @Override
    public String getParameter(String name) {
        String value = super.getParameter(name);
        String newValue = cleanXss(value);
        log.info("请求参数 value:【{}】 newValue:【{}】",value, newValue);
        return newValue;
    }

    @Override
    public String[] getParameterValues(String name) {
        String[] values = super.getParameterValues(name);
        String[] newValues = Arrays.stream(values).map(this::cleanXss).toArray(String[]::new);
        log.info("请求参数 values:【{}】 newValues:【{}】",values, newValues);
        return newValues;
    }

    /**
     * 重写参数获取方法
     *
     * @return
     */
    @Override
    public Map<String, String[]> getParameterMap() {
        Map<String, String[]> values = super.getParameterMap();
        LinkedHashMap<String, String[]> newValues = new LinkedHashMap();
        if (values != null) {
            for (String key : values.keySet()) {
                String[] valueArray = values.get(key);
                for (int i = 0; i < valueArray.length; i++) {
                    String value = valueArray[i];
                    if (!StrUtil.hasEmpty(value)) {
                        value = cleanXss(value);
                    }
                    valueArray[i] = value;
                }
                newValues.put(key, valueArray);
            }
        }
        log.info("请求参数 values:【{}】 newValues:【{}】",values, newValues);
        return newValues;
    }


    /**
     * 获取去除标签后的value值
     *
     * @param jsonObject
     */
    private void getValue(JSONObject jsonObject) {
        Set<String> keys = jsonObject.keySet();
        keys.forEach(key -> {
            Object value = jsonObject.get(key);
            if (value instanceof Collection) {
                List<?> values = (List<?>) value;
                value = values.stream().map(item -> cleanXss(item.toString())).toList();
            } else {
                value = cleanXss(value.toString());
            }
            jsonObject.put(key, value);
        });
    }

    // 自定义输入流实现
    private static class CachedServletInputStream extends ServletInputStream {
        private final ByteArrayInputStream buffer;

        public CachedServletInputStream(byte[] contents) {
            this.buffer = new ByteArrayInputStream(contents);
        }

        @Override
        public int read() {
            return buffer.read();
        }

        @Override
        public boolean isFinished() {
            return buffer.available() == 0;
        }

        @Override
        public boolean isReady() {
            return true;
        }

        @Override
        public void setReadListener(ReadListener listener) {
            throw new UnsupportedOperationException();
        }
    }

}
